// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/io/transform.h>

#include <algorithm>
#include <cstring>
#include <mutex>
#include <random>
#include <thread>
#include <utility>

#include <nebula/core/buffer.h>
#include <nebula/io/util_internal.h>

#include <turbo/utility/status.h>
#include <turbo/log/logging.h>

namespace nebula::io {

    struct TransformInputStream::Impl {
        std::shared_ptr<InputStream> wrapped_;
        TransformInputStream::TransformFunc transform_;
        std::shared_ptr<Buffer> pending_;
        int64_t pos_ = 0;
        bool closed_ = false;

        Impl(std::shared_ptr<InputStream> wrapped,
             TransformInputStream::TransformFunc transform)
                : wrapped_(std::move(wrapped)), transform_(std::move(transform)) {}

        void close() {
            closed_ = true;
            pending_.reset();
        }

        turbo::Status CheckClosed() const {
            if (closed_) {
                return turbo::invalid_argument_error("Operation on closed file");
            }
            return turbo::OkStatus();
        }
    };

    TransformInputStream::TransformInputStream(std::shared_ptr<InputStream> wrapped,
                                               TransformInputStream::TransformFunc transform)
            : impl_(new Impl{std::move(wrapped), std::move(transform)}) {}

    TransformInputStream::~TransformInputStream() {}

    turbo::Status TransformInputStream::close() {
        impl_->close();
        return impl_->wrapped_->close();
    }

    turbo::Status TransformInputStream::abort() { return impl_->wrapped_->abort(); }

    bool TransformInputStream::closed() const { return impl_->closed_; }

    turbo::Result<std::shared_ptr<Buffer>> TransformInputStream::read(int64_t nbytes) {
        TURBO_RETURN_NOT_OK(impl_->CheckClosed());

        TURBO_MOVE_OR_RAISE(auto buf, allocate_resizable_buffer(nbytes));
        TURBO_MOVE_OR_RAISE(auto bytes_read, this->read(nbytes, buf->mutable_data()));
        if (bytes_read < nbytes) {
            TURBO_RETURN_NOT_OK(buf->resize(bytes_read, /*shrink_to_fit=*/true));
        }
        return std::shared_ptr<Buffer>(std::move(buf));
    }

    turbo::Result<int64_t> TransformInputStream::read(int64_t nbytes, void *out) {
        TURBO_RETURN_NOT_OK(impl_->CheckClosed());

        if (nbytes == 0) {
            return 0;
        }

        int64_t avail_size = 0;
        std::vector<std::shared_ptr<Buffer>> avail;
        if (impl_->pending_) {
            avail.push_back(impl_->pending_);
            avail_size += impl_->pending_->size();
        }
        // Accumulate enough transformed data to satisfy read
        while (avail_size < nbytes) {
            TURBO_MOVE_OR_RAISE(auto buf, impl_->wrapped_->read(nbytes));
            const bool have_eof = (buf->size() == 0);
            // Even if EOF is met, let the transform function run a last time
            // (for example to flush internal buffers)
            TURBO_MOVE_OR_RAISE(buf, impl_->transform_(std::move(buf)));
            avail_size += buf->size();
            avail.push_back(std::move(buf));
            if (have_eof) {
                break;
            }
        }
                DKCHECK(!avail.empty());

        // Coalesce buffer data
        uint8_t *out_data = reinterpret_cast<uint8_t *>(out);
        int64_t copied_bytes = 0;
        for (size_t i = 0; i < avail.size() - 1; ++i) {
            // All buffers except the last fit fully into `nbytes`
            const auto buf = std::move(avail[i]);
                    DKCHECK_LE(buf->size(), nbytes);
            memcpy(out_data, buf->data(), static_cast<size_t>(buf->size()));
            out_data += buf->size();
            nbytes -= buf->size();
            copied_bytes += buf->size();
        }
        {
            // Last buffer: splice into `out` and `pending_`
            const auto buf = std::move(avail.back());
            const int64_t to_copy = std::min(buf->size(), nbytes);
            memcpy(out_data, buf->data(), static_cast<size_t>(to_copy));
            copied_bytes += to_copy;
            if (buf->size() > to_copy) {
                impl_->pending_ = SliceBuffer(buf, to_copy);
            } else {
                impl_->pending_.reset();
            }
        }
        impl_->pos_ += copied_bytes;
        return copied_bytes;
    }

    turbo::Result<int64_t> TransformInputStream::tell() const {
        TURBO_RETURN_NOT_OK(impl_->CheckClosed());

        return impl_->pos_;
    }

    turbo::Result<std::shared_ptr<const KeyValueMetadata>> TransformInputStream::read_metadata() {
        TURBO_RETURN_NOT_OK(impl_->CheckClosed());

        return impl_->wrapped_->read_metadata();
    }

    Future<std::shared_ptr<const KeyValueMetadata>> TransformInputStream::read_metadata_async(
            const IOContext &io_context) {
        TURBO_RETURN_NOT_OK(impl_->CheckClosed());

        return impl_->wrapped_->read_metadata_async(io_context);
    }

}  // namespace nebula::io

